/*
 * Copyright (C) 2011-2021 Intel Corporation.
 * SPDX-License-Identifier: MIT
 */

#include "pin.H"
#include <iostream>
#include <cstdlib>
#include <signal.h>
#include <list>
#include <string>
using std::hex;

using std::cerr;
using std::endl;
using std::list;
using std::string;

extern "C" void ToolRaiseAccessInvalidAddressException(ADDRINT** addr, ADDRINT val);
extern "C" ADDRINT ToolCatchAccessInvalidAddressException;
extern "C" ADDRINT ToolIpAccessInvalidAddressException;

extern "C" ADDRINT ToolRaiseIntDivideByZeroException(ADDRINT catchAddr, ADDRINT code);
extern "C" ADDRINT ToolCatchIntDivideByZeroException;
extern "C" ADDRINT ToolIpIntDivideByZeroException;

enum ToolExceptionEnumeration
{
    TEE_ACCESS_INVALID_ADDRESS,
    TEE_INT_DIV_ZERO
};

struct ToolException
{
    ADDRINT ip;
    ADDRINT fixIp;
    string name;
    bool* handled;
};

bool handledPinException             = false;
bool handledToolAccessInvalidAddress = false;
bool handledToolIntDivZero           = false;
bool handledApplicationDivZero       = false;

ToolException toolExceptions[2];
list< ADDRINT > inslist;

// This is the tool's generic handler. It handles the excpetion defined in toolException.
static EXCEPT_HANDLING_RESULT toolHandler(ToolException* toolException, EXCEPTION_INFO* exptInfo, PHYSICAL_CONTEXT* pctxt)
{
    ADDRINT exptAddr = PIN_GetExceptionAddress(exptInfo);
    ADDRINT ip       = PIN_GetPhysicalContextReg(pctxt, REG_INST_PTR);
    if (ip == toolException->ip)
    { // Only handle this exception.
        cerr << "TOOL: Identified " << toolException->name << " exception in the tool, fixing the problem..." << endl;
        PIN_SetPhysicalContextReg(pctxt, REG_INST_PTR, toolException->fixIp);
        if (exptAddr != ip)
        {
            cerr << hex << "TOOL ERROR: Exception address 0x" << exptAddr << " does not match ip 0x" << ip << endl;
        }
        else
        {
            *(toolException->handled) = true;
        }
        return EHR_HANDLED;
    }
    return EHR_UNHANDLED;
}

// Global handler for SIGFPE exceptions.
// This handler only deals with a specific int-div-zero exception, the one raised in "generateToolException()".
// All other exceptions are not handled here and are forwarded to the next level.
static EXCEPT_HANDLING_RESULT onException(THREADID tid, EXCEPTION_INFO* exptInfo, PHYSICAL_CONTEXT* pctxt, VOID* v)
{
    return toolHandler(&toolExceptions[TEE_INT_DIV_ZERO], exptInfo, pctxt);
}

// Thread specific handler for PIN_Try exceptions.
// This handler only deals with a specific access-invalid-address exception, the one raised in "generateToolException()".
// All other exceptions are not handled here and are forwarded to the next level.
static EXCEPT_HANDLING_RESULT onTryException(THREADID tid, EXCEPTION_INFO* exptInfo, PHYSICAL_CONTEXT* pctxt, VOID* v)
{
    return toolHandler(&toolExceptions[TEE_ACCESS_INVALID_ADDRESS], exptInfo, pctxt);
}

// Pin generates an access violation in SafeCopy (due to the given arguments).
// This should return an exception with a null exception address
static void generatePinException(ADDRINT address)
{
    cerr << "TOOL: Generate pin exception" << endl;
    VOID* to      = (VOID*)address;
    ADDRINT* from = 0;
    EXCEPTION_INFO exptInfo;
    size_t toCopy = 1024;
    if (PIN_SafeCopyEx(to, from, toCopy, &exptInfo) != toCopy)
    {
        if (PIN_GetExceptionCode(&exptInfo) != EXCEPTCODE_ACCESS_INVALID_ADDRESS)
        {
            cerr << "TOOL ERROR: PIN_SafeCopyEx returned with an unexpected exception code" << endl;
            cerr << PIN_ExceptionToString(&exptInfo) << endl;
            return;
        }
        ADDRINT exptAddr = PIN_GetExceptionAddress(&exptInfo);
        if (exptAddr != 0)
        {
            cerr << "TOOL ERROR: Exception from PIN_SafeCopyEx returned with non NULL address: 0x" << hex << exptAddr << endl;
            return;
        }
        else
        {
            cerr << "TOOL: PIN_SafeCopyEx failed as expected, continue test..." << endl;
            handledPinException = true;
        }
    }
}

// Tool generates a div zero exception.
// This should return an exception address which equals the IP value.
// The address is in the analysis routine placed in the codecache.
static void generateToolException()
{
    ToolRaiseIntDivideByZeroException(ToolCatchIntDivideByZeroException, 0);
    PIN_TryStart(PIN_ThreadId(), onTryException, 0);
    ADDRINT* addressPair[2];
    addressPair[0] = 0;
    addressPair[1] = (ADDRINT*)malloc(sizeof(ADDRINT));
    ToolRaiseAccessInvalidAddressException(addressPair, 0x12345);
    PIN_TryEnd(PIN_ThreadId());
}

// Catch a SIGFPE signal generated by a div zero in the application.
// This should return an exception address in the application's native code, not the instrumented code in the codecache.
// Fix the problem by advancing the IP to the next instruction - advance the IP in the context (native code) and let Pin
// translate to the actual address in the codecache.
static bool onSIGFPE(THREADID tid, INT32 sig, CONTEXT* ctxt, BOOL hasHandler, const EXCEPTION_INFO* exptInfo, VOID* v)
{
    ADDRINT exptAddr                    = PIN_GetExceptionAddress(exptInfo);
    ADDRINT ip                          = PIN_GetContextReg(ctxt, REG_INST_PTR);
    list< ADDRINT >::const_iterator it  = inslist.begin();
    list< ADDRINT >::const_iterator end = inslist.end();
    for (; it != end; ++it)
    {
        if (*it == ip)
        {
            cerr << "TOOL: Identified div zero instruction in the application, fixing the problem..." << endl;
            ++it;                                       // Advance the iterator to the next instruction.
            PIN_SetContextReg(ctxt, REG_INST_PTR, *it); // Load the new IP.
            if (exptAddr == ip)
            {
                handledApplicationDivZero = true;
            }
            else
            {
                cerr << hex << "TOOL ERROR: Exception address 0x" << exptAddr << " does not match ip 0x" << ip << endl;
            }
            return false;
        }
    }
    return true;
}

// Instrument the application functions - intentionally generate exceptions in the proper places and build the instList.
static VOID onRoutine(RTN rtn, VOID*)
{
    if (RTN_Name(rtn).find("pinException") != string::npos)
    {
        RTN_Open(rtn);
        RTN_InsertCall(rtn, IPOINT_AFTER, AFUNPTR(generatePinException), IARG_ADDRINT, RTN_Address(rtn), IARG_END);
        RTN_Close(rtn);
    }

    else if (RTN_Name(rtn).find("toolException") != string::npos)
    {
        RTN_Open(rtn);
        RTN_InsertCall(rtn, IPOINT_AFTER, AFUNPTR(generateToolException), IARG_END);
        RTN_Close(rtn);
    }

    else if (RTN_Name(rtn).find("appException") != string::npos)
    {
        RTN_Open(rtn);

        // Build a list of all the instrucions in the routine.
        // The tool will use this to fix the div zero exception by advancing the IP to the next instruction.
        for (INS ins = RTN_InsHead(rtn); INS_Valid(ins); ins = INS_Next(ins))
        {
            inslist.push_back(INS_Address(ins));
        }

        RTN_Close(rtn);
    }
}

// Check that all four scenarios were handled.
VOID Fini(INT32 code, VOID* v)
{
    bool failed = false;
    if (!handledPinException)
    {
        cerr << "TOOL ERROR: PIN_SafeCopy exception was not handled properly in the test." << endl;
        failed = true;
    }
    if (!handledToolAccessInvalidAddress)
    {
        cerr << "TOOL ERROR: Tool access invalid address exception was not handled properly in the test." << endl;
        failed = true;
    }
    if (!handledToolIntDivZero)
    {
        cerr << "TOOL ERROR: Tool int div zero exception was not handled properly in the test." << endl;
        failed = true;
    }
    if (!handledApplicationDivZero)
    {
        cerr << "TOOL ERROR: Application div zero exception was not handled properly in the test." << endl;
        failed = true;
    }
    if (failed)
    {
        exit(-1);
    }
    else
    {
        cerr << "TOOL: Test completed successfully!" << endl;
    }
}

// Tool's main function
int main(int argc, char* argv[])
{
    PIN_InitSymbols();
    PIN_Init(argc, argv);

    toolExceptions[TEE_ACCESS_INVALID_ADDRESS].ip      = reinterpret_cast< ADDRINT >(&ToolIpAccessInvalidAddressException);
    toolExceptions[TEE_ACCESS_INVALID_ADDRESS].fixIp   = reinterpret_cast< ADDRINT >(&ToolCatchAccessInvalidAddressException);
    toolExceptions[TEE_ACCESS_INVALID_ADDRESS].name    = "access invalid address";
    toolExceptions[TEE_ACCESS_INVALID_ADDRESS].handled = &handledToolAccessInvalidAddress;

    toolExceptions[TEE_INT_DIV_ZERO].ip      = reinterpret_cast< ADDRINT >(&ToolIpIntDivideByZeroException);
    toolExceptions[TEE_INT_DIV_ZERO].fixIp   = reinterpret_cast< ADDRINT >(&ToolCatchIntDivideByZeroException);
    toolExceptions[TEE_INT_DIV_ZERO].name    = "int div zero";
    toolExceptions[TEE_INT_DIV_ZERO].handled = &handledToolIntDivZero;

    PIN_AddInternalExceptionHandler(onException, 0);
    RTN_AddInstrumentFunction(onRoutine, 0);
    PIN_InterceptSignal(SIGFPE, onSIGFPE, 0);
    PIN_AddFiniFunction(Fini, 0);

    PIN_StartProgram();
    return 0;
}
